import os
import re
import json
import csv
from PyPDF2 import PdfReader
from collections import defaultdict

# 配置变量 - 修改这些变量以适应你的需求
PDF_FOLDER = "./pdf_files"  # PDF文件所在的文件夹路径
OUTPUT_DIR = "./ones"        # 输出目录
OUTPUT_CSV = os.path.join(OUTPUT_DIR, "pdf_extracted_info.csv")  # CSV输出文件
OUTPUT_JSON = os.path.join(OUTPUT_DIR, "pdf_extracted_info.json")  # JSON输出文件

def extract_largest_font_text(pdf_path):
    """
    提取PDF中字体最大的文本作为标题
    """
    with open(pdf_path, 'rb') as file:
        reader = PdfReader(file)
        first_page = reader.pages[0]
        
        # 分析字体大小和对应的文本
        font_sizes = defaultdict(list)
        
        # 提取页面内容
        content = first_page.extract_text()
        if not content:
            return os.path.basename(pdf_path)
        
        # 简单的字体大小分析（基于PyPDF2的局限性）
        # 注意：PyPDF2对字体大小的支持有限，可能需要更专业的库如pdfminer.six
        # 这里我们使用简单的启发式方法：假设标题通常是第一段且较短
        lines = content.split('\n')
        if lines:
            # 取第一行作为标题候选
            return lines[0].strip()
        
        return os.path.basename(pdf_path)

def extract_pdf_info(pdf_path):
    """
    从PDF文件中提取标题和GitHub链接
    """
    title = extract_largest_font_text(pdf_path)
    
    with open(pdf_path, 'rb') as file:
        reader = PdfReader(file)
        
        # 查找GitHub链接
        github_links = set()
        for page in reader.pages:
            page_text = page.extract_text()
            if page_text:
                # 匹配GitHub链接的正则表达式
                matches = re.findall(r'https?://github\.com/[^\s)\]]+', page_text)
                github_links.update(matches)
        
        return {
            'file_name': os.path.basename(pdf_path),
            'title': title,
            'github_links': list(github_links)
        }

def process_pdf_folder(folder_path):
    """
    处理文件夹中的所有PDF文件
    """
    results = []
    
    for filename in os.listdir(folder_path):
        if filename.lower().endswith('.pdf'):
            pdf_path = os.path.join(folder_path, filename)
            try:
                result = extract_pdf_info(pdf_path)
                results.append(result)
                print(f"处理完成: {filename}")
            except Exception as e:
                print(f"处理 {filename} 时出错: {str(e)}")
    
    return results

def save_to_csv(results, output_file):
    """
    将结果保存为CSV文件
    """
    with open(output_file, 'w', newline='', encoding='utf-8') as csvfile:
        fieldnames = ['file_name', 'title', 'github_links']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        
        writer.writeheader()
        for result in results:
            # 将GitHub链接列表转换为字符串
            result_csv = result.copy()
            result_csv['github_links'] = ', '.join(result['github_links'])
            writer.writerow(result_csv)

def save_to_json(results, output_file):
    """
    将结果保存为JSON文件
    """
    with open(output_file, 'w', encoding='utf-8') as jsonfile:
        json.dump(results, jsonfile, ensure_ascii=False, indent=4)

def main():
    # 验证文件夹是否存在
    if not os.path.isdir(PDF_FOLDER):
        print(f"错误: 指定的文件夹 '{PDF_FOLDER}' 不存在")
        return
    
    # 处理PDF文件
    print(f"正在处理文件夹: {PDF_FOLDER}")
    results = process_pdf_folder(PDF_FOLDER)
    
    if not results:
        print("没有找到可处理的PDF文件")
        return
    
    # 保存结果
    save_to_csv(results, OUTPUT_CSV)
    save_to_json(results, OUTPUT_JSON)
    
    print(f"\n处理完成! 结果已保存到:")
    print(f"- CSV文件: {OUTPUT_CSV}")
    print(f"- JSON文件: {OUTPUT_JSON}")

if __name__ == "__main__":
    main()